{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "FH3IRpYTta2v"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FH3IRpYTta2v"
      },
      "source": [
        "##### Copyright 2023 The IREE Authors"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mWGa71_Ct2ug",
        "cellView": "form"
      },
      "source": [
        "#@title Licensed under the Apache License v2.0 with LLVM Exceptions.\n",
        "# See https://llvm.org/LICENSE.txt for license information.\n",
        "# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h5s6ncerSpc5"
      },
      "source": [
        "# Dynamic Shapes\n",
        "\n",
        "This notebook\n",
        "\n",
        "1. Creates a PyTorch program with dynamic shapes using [SHARK-Turbine](https://github.com/nod-ai/SHARK-Turbine)'s advanced AOT toolkit\n",
        "2. Compiles the program to an IREE VM bytecode module\n",
        "3. Tests running the compiled VM module using IREE's runtime\n",
        "4. Downloads compilation artifacts for use with the native (C API) sample application"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s2bScbYkP6VZ",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7b268798-a20d-4df4-f00d-ed7811f77767"
      },
      "source": [
        "#@title General setup\n",
        "\n",
        "import os\n",
        "import tempfile\n",
        "\n",
        "ARTIFACTS_DIR = os.path.join(tempfile.gettempdir(), \"iree\", \"colab_artifacts\")\n",
        "os.makedirs(ARTIFACTS_DIR, exist_ok=True)\n",
        "print(f\"Using artifacts directory '{ARTIFACTS_DIR}'\")"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using artifacts directory '/tmp/iree/colab_artifacts'\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "#@title Uninstall existing packages\n",
        "#   This avoids some warnings when installing specific PyTorch packages below.\n",
        "!python -m pip uninstall -y fastai torchaudio torchdata torchtext torchvision"
      ],
      "metadata": {
        "id": "y9KOsqosg6Ms"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Install SHARK-Turbine\n",
        "\n",
        "# Limit cell height.\n",
        "from IPython.display import Javascript\n",
        "display(Javascript('''google.colab.output.setIframeHeight(0, true, {maxHeight: 300})'''))\n",
        "\n",
        "!python -m pip install shark-turbine"
      ],
      "metadata": {
        "id": "SdCAvI3sqBO7",
        "outputId": "2be248d9-bf6b-475e-c44a-aa529f20de23",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 300
        }
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "google.colab.output.setIframeHeight(0, true, {maxHeight: 300})"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting shark-turbine\n",
            "  Downloading shark-turbine-0.9.1.dev3.tar.gz (60 kB)\n",
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/60.2 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.2/60.2 kB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from shark-turbine) (1.23.5)\n",
            "Collecting iree-compiler>=20231004.665 (from shark-turbine)\n",
            "  Downloading iree_compiler-20231004.665-cp310-cp310-manylinux_2_28_x86_64.whl (57.2 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m57.2/57.2 MB\u001b[0m \u001b[31m17.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting iree-runtime>=20231004.665 (from shark-turbine)\n",
            "  Downloading iree_runtime-20231004.665-cp310-cp310-manylinux_2_28_x86_64.whl (7.8 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m91.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: torch>=2.1.0 in /usr/local/lib/python3.10/dist-packages (from shark-turbine) (2.1.0+cu118)\n",
            "Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from iree-compiler>=20231004.665->shark-turbine) (6.0.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->shark-turbine) (3.12.4)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->shark-turbine) (4.5.0)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->shark-turbine) (1.12)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->shark-turbine) (3.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->shark-turbine) (3.1.2)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->shark-turbine) (2023.6.0)\n",
            "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.10/dist-packages (from torch>=2.1.0->shark-turbine) (2.1.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=2.1.0->shark-turbine) (2.1.3)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=2.1.0->shark-turbine) (1.3.0)\n",
            "Building wheels for collected packages: shark-turbine\n",
            "  Building wheel for shark-turbine (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for shark-turbine: filename=shark_turbine-0.9.1.dev3-py3-none-any.whl size=70102 sha256=507dec827b9a2eea18f47c6ebdc84347c9956b8f2e0b186d3107a006e0742d81\n",
            "  Stored in directory: /root/.cache/pip/wheels/e9/78/0f/88c9d8224ef1550fe00b18a014eab5121f26264e2261f31926\n",
            "Successfully built shark-turbine\n",
            "Installing collected packages: iree-runtime, iree-compiler, shark-turbine\n",
            "Successfully installed iree-compiler-20231004.665 iree-runtime-20231004.665 shark-turbine-0.9.1.dev3\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Report version information\n",
        "!echo \"Installed SHARK-Turbine, $(python -m pip show shark_turbine | grep Version)\"\n",
        "\n",
        "!echo -e \"\\nInstalled IREE, compiler version information:\"\n",
        "!iree-compile --version\n",
        "\n",
        "import torch\n",
        "print(\"\\nInstalled PyTorch, version:\", torch.__version__)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Oj5I6R9LI7t_",
        "outputId": "35d79e6a-7bd0-46e1-8113-5af1a7bcbb5b"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Installed SHARK-Turbine, Version: 0.9.1.dev3\n",
            "\n",
            "Installed IREE, compiler version information:\n",
            "IREE (https://iree.dev):\n",
            "  IREE compiler version 20231004.665 @ bb51f6f1a1b4ee619fb09a7396f449dadb211447\n",
            "  LLVM version 18.0.0git\n",
            "  Optimized build\n",
            "\n",
            "Installed PyTorch, version: 2.1.0+cu118\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Create a program using PyTorch + SHARK-Turbine\n",
        "\n",
        "NOTE: as in other domains, providing more information to a compiler allows it\n",
        "to generate more efficient code. As a general rule, the slowest varying\n",
        "dimensions of program data like batch index or timestep are safer to treat as\n",
        "dynamic than faster varying dimensions like image x/y/channel. See\n",
        "[this paper](https://arxiv.org/pdf/2006.03031.pdf) for a discussion of the\n",
        "challenges imposed by dynamic shapes and one project's approach to addressing\n",
        "them."
      ],
      "metadata": {
        "id": "C3mhaullI940"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Define a sample `shark_turbine.aot.CompiledModule` using dynamic shapes\n",
        "\n",
        "import shark_turbine.aot as aot\n",
        "\n",
        "class DynamicShapesModule(aot.CompiledModule, export_name=\"module\"):\n",
        "  # reduce_sum_1d (dynamic input size, static output size)\n",
        "  #   tensor<?xi32> -> tensor<i32>\n",
        "  #   e.g. [1, 2, 3] -> 6\n",
        "  def reduce_sum_1d(self, values=aot.AbstractTensor(None, dtype=torch.int32)):\n",
        "    return self.compute_reduce_sum_1d(values)\n",
        "\n",
        "  @aot.jittable\n",
        "  def compute_reduce_sum_1d(values):\n",
        "    return torch.sum(values, dtype=torch.int32)\n",
        "\n",
        "  # reduce_sum_2d (partially dynamic input size, static output size)\n",
        "  #   tensor<?x3xi32> -> tensor<3xi32>\n",
        "  #   e.g. [[1, 2, 3], [10, 20, 30]] -> [11, 22, 33]\n",
        "  def reduce_sum_2d(self, values=aot.AbstractTensor(None, 3, dtype=torch.int32)):\n",
        "    return self.compute_reduce_sum_2d(values)\n",
        "\n",
        "  @aot.jittable\n",
        "  def compute_reduce_sum_2d(values):\n",
        "    return torch.sum(values, 0, dtype=torch.int32)\n",
        "\n",
        "  # add_one (dynamic input size, dynamic output size)\n",
        "  #   tensor<?xi32>) -> tensor<?xi32>\n",
        "  #   e.g. [1, 2, 3] -> [2, 3, 4]\n",
        "  def add_one(self, values=aot.AbstractTensor(None, dtype=torch.int32)):\n",
        "    return self.compute_add_one(values)\n",
        "\n",
        "  @aot.jittable\n",
        "  def compute_add_one(values):\n",
        "    return values + 1"
      ],
      "metadata": {
        "id": "vsf9F4WxI_DX"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from iree.compiler.ir import Context\n",
        "\n",
        "# Import into MLIR and save to disk.\n",
        "dynamic_shapes_instance = DynamicShapesModule(context=Context())\n",
        "imported_mlir_path = os.path.join(ARTIFACTS_DIR, \"dynamic_shapes.mlir\")\n",
        "aot.CompiledModule.save_mlir(dynamic_shapes_instance, imported_mlir_path)\n",
        "print(f\"Wrote MLIR to path '{imported_mlir_path}'\")\n",
        "\n",
        "# Inspect the IR.\n",
        "# Note the question marks for dynamic shapes in types, like `tensor<?xi32>`.\n",
        "print(\"\\nDynamic Shapes MLIR:\")\n",
        "!cat {imported_mlir_path}"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_OQIpOtNr4Gh",
        "outputId": "888c0bf3-bec6-403c-9993-ad45d21364fb"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Wrote MLIR to path '/tmp/iree/colab_artifacts/dynamic_shapes.mlir'\n",
            "\n",
            "Dynamic Shapes MLIR:\n",
            "#map = affine_map<(d0) -> (d0)>\n",
            "#map1 = affine_map<(d0) -> ()>\n",
            "#map2 = affine_map<(d0, d1) -> (d0, d1)>\n",
            "#map3 = affine_map<(d0, d1) -> (d1)>\n",
            "module @module {\n",
            "  func.func @reduce_sum_1d(%arg0: tensor<?xi32>) -> tensor<i32> attributes {torch.args_schema = \"[1, {\\22type\\22: \\22builtins.tuple\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: \\22builtins.list\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]}, {\\22type\\22: \\22builtins.dict\\22, \\22context\\22: \\22[]\\22, \\22children_spec\\22: []}]}]\", torch.return_schema = \"[1, {\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]\"} {\n",
            "    %0 = call @compute_reduce_sum_1d(%arg0) : (tensor<?xi32>) -> tensor<i32>\n",
            "    return %0 : tensor<i32>\n",
            "  }\n",
            "  func.func private @compute_reduce_sum_1d(%arg0: tensor<?xi32>) -> tensor<i32> {\n",
            "    %c0_i32 = arith.constant 0 : i32\n",
            "    %0 = tensor.empty() : tensor<i32>\n",
            "    %1 = linalg.fill ins(%c0_i32 : i32) outs(%0 : tensor<i32>) -> tensor<i32>\n",
            "    %2 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = [\"reduction\"]} ins(%arg0 : tensor<?xi32>) outs(%1 : tensor<i32>) {\n",
            "    ^bb0(%in: i32, %out: i32):\n",
            "      %3 = arith.addi %in, %out : i32\n",
            "      linalg.yield %3 : i32\n",
            "    } -> tensor<i32>\n",
            "    return %2 : tensor<i32>\n",
            "  }\n",
            "  func.func @reduce_sum_2d(%arg0: tensor<?x3xi32>) -> tensor<3xi32> attributes {torch.args_schema = \"[1, {\\22type\\22: \\22builtins.tuple\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: \\22builtins.list\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]}, {\\22type\\22: \\22builtins.dict\\22, \\22context\\22: \\22[]\\22, \\22children_spec\\22: []}]}]\", torch.return_schema = \"[1, {\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]\"} {\n",
            "    %0 = call @compute_reduce_sum_2d(%arg0) : (tensor<?x3xi32>) -> tensor<3xi32>\n",
            "    return %0 : tensor<3xi32>\n",
            "  }\n",
            "  func.func private @compute_reduce_sum_2d(%arg0: tensor<?x3xi32>) -> tensor<3xi32> {\n",
            "    %c0_i32 = arith.constant 0 : i32\n",
            "    %0 = tensor.empty() : tensor<3xi32>\n",
            "    %1 = linalg.fill ins(%c0_i32 : i32) outs(%0 : tensor<3xi32>) -> tensor<3xi32>\n",
            "    %2 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = [\"reduction\", \"parallel\"]} ins(%arg0 : tensor<?x3xi32>) outs(%1 : tensor<3xi32>) {\n",
            "    ^bb0(%in: i32, %out: i32):\n",
            "      %3 = arith.addi %in, %out : i32\n",
            "      linalg.yield %3 : i32\n",
            "    } -> tensor<3xi32>\n",
            "    return %2 : tensor<3xi32>\n",
            "  }\n",
            "  func.func @add_one(%arg0: tensor<?xi32>) -> tensor<?xi32> attributes {torch.args_schema = \"[1, {\\22type\\22: \\22builtins.tuple\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: \\22builtins.list\\22, \\22context\\22: \\22null\\22, \\22children_spec\\22: [{\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]}, {\\22type\\22: \\22builtins.dict\\22, \\22context\\22: \\22[]\\22, \\22children_spec\\22: []}]}]\", torch.return_schema = \"[1, {\\22type\\22: null, \\22context\\22: null, \\22children_spec\\22: []}]\"} {\n",
            "    %0 = call @compute_add_one(%arg0) : (tensor<?xi32>) -> tensor<?xi32>\n",
            "    return %0 : tensor<?xi32>\n",
            "  }\n",
            "  func.func private @compute_add_one(%arg0: tensor<?xi32>) -> tensor<?xi32> {\n",
            "    %c0 = arith.constant 0 : index\n",
            "    %c1_i32 = arith.constant 1 : i32\n",
            "    %dim = tensor.dim %arg0, %c0 : tensor<?xi32>\n",
            "    %0 = tensor.empty(%dim) : tensor<?xi32>\n",
            "    %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = [\"parallel\"]} ins(%arg0 : tensor<?xi32>) outs(%0 : tensor<?xi32>) {\n",
            "    ^bb0(%in: i32, %out: i32):\n",
            "      %2 = arith.addi %in, %c1_i32 : i32\n",
            "      linalg.yield %2 : i32\n",
            "    } -> tensor<?xi32>\n",
            "    return %1 : tensor<?xi32>\n",
            "  }\n",
            "}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Test the imported program\n",
        "\n",
        "_Note: you can stop after each step and use intermediate outputs with other tools outside of Colab._\n",
        "\n",
        "_See the [README](https://github.com/openxla/iree/tree/main/samples/dynamic_shapes#instructions) for more details and example command line instructions._\n",
        "\n",
        "* _The \"imported MLIR\" (above) can be used by IREE's generic compiler tools_\n",
        "* _The \"binary\" can be saved and used by runtime applications_\n",
        "\n",
        "_The specific point at which you switch from Python to native tools will depend on your project._"
      ],
      "metadata": {
        "id": "z6w_Pbl6tUtJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Export and compile.\n",
        "exported_output = aot.export(DynamicShapesModule)\n",
        "\n",
        "# Compile to a file on disk for usage outside of Python.\n",
        "flatbuffer_path = os.path.join(ARTIFACTS_DIR, \"dynamic_shapes_cpu.vmfb\")\n",
        "exported_output.compile(save_to=flatbuffer_path)\n",
        "print(f\"Wrote compiled program to path '{flatbuffer_path}'\")\n",
        "\n",
        "# Compile into memory for testing.\n",
        "binary = exported_output.compile(save_to=None)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0PGyH1tvI_Ic",
        "outputId": "23b53928-4d77-461f-e4b8-b2c8ffb25ef0"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Wrote compiled program to path '/tmp/iree/colab_artifacts/dynamic_shapes_cpu.vmfb'\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import iree.runtime as ireert\n",
        "import numpy as np\n",
        "\n",
        "# Use the IREE runtime API to test the compiled program.\n",
        "config = ireert.Config(\"local-task\")\n",
        "vm_module = ireert.load_vm_module(\n",
        "    ireert.VmModule.wrap_buffer(config.vm_instance, binary.map_memory()),\n",
        "    config,\n",
        ")\n",
        "\n",
        "print(vm_module.reduce_sum_1d(np.array([1, 10, 100], dtype=np.int32)).to_host())\n",
        "print(vm_module.reduce_sum_2d(np.array([[1, 2, 3], [10, 20, 30]], dtype=np.int32)).to_host())\n",
        "print(vm_module.reduce_sum_2d(np.array([[1, 2, 3], [10, 20, 30], [100, 200, 300]], dtype=np.int32)).to_host())\n",
        "print(vm_module.add_one(np.array([1, 10, 100], dtype=np.int32)).to_host())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9ilJY15BI_LD",
        "outputId": "57db6e52-83f1-4283-fc08-31e743cc9b42"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "111\n",
            "[11 22 33]\n",
            "[111 222 333]\n",
            "[  2  11 101]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Download compilation artifacts"
      ],
      "metadata": {
        "id": "3mizlpY9uJEW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ARTIFACTS_ZIP = \"/tmp/dynamic_shapes_colab_artifacts.zip\"\n",
        "\n",
        "print(f\"Zipping '{ARTIFACTS_DIR}' to '{ARTIFACTS_ZIP}' for download...\")\n",
        "!cd {ARTIFACTS_DIR} && zip -r {ARTIFACTS_ZIP} .\n",
        "\n",
        "# Note: you can also download files using Colab's file explorer\n",
        "try:\n",
        "  from google.colab import files\n",
        "  print(\"Downloading the artifacts zip file...\")\n",
        "  files.download(ARTIFACTS_ZIP)\n",
        "except ImportError:\n",
        "  print(\"Missing google_colab Python package, can't download files\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 86
        },
        "id": "dgaXpdiWuGtx",
        "outputId": "dc0fbca1-c5b0-44f9-e1ff-9bf1307c049f"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Zipping '/tmp/iree/colab_artifacts' to '/tmp/dynamic_shapes_colab_artifacts.zip' for download...\n",
            "  adding: dynamic_shapes_cpu.vmfb (deflated 66%)\n",
            "  adding: dynamic_shapes.mlir (deflated 82%)\n",
            "Downloading the artifacts zip file...\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "\n",
              "    async function download(id, filename, size) {\n",
              "      if (!google.colab.kernel.accessAllowed) {\n",
              "        return;\n",
              "      }\n",
              "      const div = document.createElement('div');\n",
              "      const label = document.createElement('label');\n",
              "      label.textContent = `Downloading \"${filename}\": `;\n",
              "      div.appendChild(label);\n",
              "      const progress = document.createElement('progress');\n",
              "      progress.max = size;\n",
              "      div.appendChild(progress);\n",
              "      document.body.appendChild(div);\n",
              "\n",
              "      const buffers = [];\n",
              "      let downloaded = 0;\n",
              "\n",
              "      const channel = await google.colab.kernel.comms.open(id);\n",
              "      // Send a message to notify the kernel that we're ready.\n",
              "      channel.send({})\n",
              "\n",
              "      for await (const message of channel.messages) {\n",
              "        // Send a message to notify the kernel that we're ready.\n",
              "        channel.send({})\n",
              "        if (message.buffers) {\n",
              "          for (const buffer of message.buffers) {\n",
              "            buffers.push(buffer);\n",
              "            downloaded += buffer.byteLength;\n",
              "            progress.value = downloaded;\n",
              "          }\n",
              "        }\n",
              "      }\n",
              "      const blob = new Blob(buffers, {type: 'application/binary'});\n",
              "      const a = document.createElement('a');\n",
              "      a.href = window.URL.createObjectURL(blob);\n",
              "      a.download = filename;\n",
              "      div.appendChild(a);\n",
              "      a.click();\n",
              "      div.remove();\n",
              "    }\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Javascript object>"
            ],
            "application/javascript": [
              "download(\"download_e2630f9b-e811-4164-b2d8-80cf52f17145\", \"dynamic_shapes_colab_artifacts.zip\", 5699)"
            ]
          },
          "metadata": {}
        }
      ]
    }
  ]
}